import tensorflow as tf
from lstm import BNLSTMCell, LSTMStateTuple
from initializer import get_zero_shape, get_normal_shape

def data_type():
    return tf.float32

# Configuration of Model
class Config(object):
    init_scale = 0.1     # scale for init of W
    learning_rate = 1.0  # lr
    max_grad_norm = 5    # grad clip norm
    num_layers = 2       # num of Layer
    num_steps = 20       # T
    hidden_size = 200    # D
    max_epoch = 4        # epoch after which start lr decay
    max_max_epoch = 13   # training epoch
    keep_prob = 1.0
    lr_decay = 0.5
    batch_size = 20      # B
    vocab_size = 10000   # vocabs in training article
def get_config():
    return Config()

# Input info of feeding data
class PTBInput(object):
    def __init__(self, config, data, name=None):
        self.batch_size = batch_size = config.batch_size
        self.num_steps = num_steps = config.num_steps
        self.epoch_size = ((len(data) // batch_size) - 1) // num_steps
        # shape: (batch_size, num_steps)
        # targets is next batch data
        self.input_data, self.targets = _producer(
            data, batch_size, num_steps, name=name)

def _producer(raw_data, batch_size, num_steps, name=None):
    """Iterate on the raw PTB data.
    This chunks up raw_data into batches of examples and returns Tensors that
    are drawn from these batches.
    Args:
        raw_data: one of the raw data outputs from ptb_raw_data.
        batch_size: int, the batch size.
        num_steps: int, the number of unrolls.
        name: the name of this operation (optional).
    Returns:
        A pair of Tensors, each shaped [batch_size, num_steps]. The second element
        of the tuple is the same data time-shifted to the right by one.
    Raises:
        tf.errors.InvalidArgumentError: if batch_size or num_steps are too high.
    """
    with tf.name_scope(name, "PTBProducer", [raw_data, batch_size, num_steps]):
        raw_data = tf.convert_to_tensor(raw_data, name="raw_data", dtype=tf.int32)

        data_len = tf.size(raw_data)
        batch_len = data_len // batch_size
        data = tf.reshape(raw_data[0 : batch_size * batch_len],
                                            [batch_size, batch_len])
        # The last batch is only used for prediction
        epoch_size = (batch_len - 1) // num_steps
        assertion = tf.assert_positive(
                epoch_size,
                message="epoch_size == 0, decrease batch_size or num_steps")
        with tf.control_dependencies([assertion]):
            epoch_size = tf.identity(epoch_size, name="epoch_size")

        i = tf.train.range_input_producer(epoch_size, shuffle=False).dequeue()
        x = tf.strided_slice(data, [0, i * num_steps],
                                                 [batch_size, (i + 1) * num_steps])
        x.set_shape([batch_size, num_steps])
        y = tf.strided_slice(data, [0, i * num_steps + 1],
                                                 [batch_size, (i + 1) * num_steps + 1])
        y.set_shape([batch_size, num_steps])
        return x, y

class PTBModel(object):
    """The PTB model."""
    def __init__(self, is_training, config, input_, mode):
        self._input = input_

        batch_size = input_.batch_size
        num_steps = input_.num_steps
        size = config.hidden_size
        vocab_size = config.vocab_size

        self._training = is_training

        # Slightly better results can be obtained with forget gate biases
        # initialized to 1 but the hyperparameters of the model would need to be
        # different than reported in the paper.
        def lstm_cell():
            # TODO 1 of 2
            # return tf.contrib.rnn.BasicLSTMCell(
            #             size, forget_bias=0.0, state_is_tuple=True,
            #             reuse=tf.get_variable_scope().reuse)
            return BNLSTMCell(
                        size, num_steps, is_training,
                        mode=mode, forget_bias=0.0)

        attn_cell = lstm_cell
        if is_training and config.keep_prob < 1:
            def attn_cell():
                return tf.contrib.rnn.DropoutWrapper(
                        lstm_cell(), output_keep_prob=config.keep_prob)
        cell = tf.contrib.rnn.MultiRNNCell(
                [attn_cell() for _ in range(config.num_layers)], state_is_tuple=True)

        # TODO 2 of 2
        # self._initial_state = cell.zero_state(batch_size, data_type())
        self._initial_state = tuple(
            LSTMStateTuple(tf.constant(0, tf.int32),
                get_zero_shape(size, batch_size, data_type()),
                get_normal_shape(size, batch_size, data_type()))
            for i in range(config.num_layers))

        with tf.device("/cpu:0"):
            embedding = tf.get_variable(
                    "embedding", [vocab_size, size], dtype=data_type())
            inputs = tf.nn.embedding_lookup(embedding, input_.input_data)
        # inputs shape: (B, T, D)

        if is_training and config.keep_prob < 1:
            inputs = tf.nn.dropout(inputs, config.keep_prob)

        # Simplified version of models/tutorials/rnn/rnn.py's rnn().
        # This builds an unrolled LSTM for tutorial purposes only.
        # In general, use the rnn() or state_saving_rnn() from rnn.py.
        #
        # The alternative version of the code below is:
        #
        # inputs = tf.unstack(inputs, num=num_steps, axis=1)
        # outputs, state = tf.contrib.rnn.static_rnn(
        #     cell, inputs, initial_state=self._initial_state)


        outputs, state = tf.nn.dynamic_rnn(cell, inputs,
            initial_state=self._initial_state, dtype=tf.float32)
        output = tf.reshape(outputs, [-1, size])



        # outputs = []
        # state = self._initial_state
        # with tf.variable_scope("RNN"):
        #     for time_step in range(num_steps):
        #         if time_step > 0: tf.get_variable_scope().reuse_variables()
        #         (cell_output, state) = cell(inputs[:, time_step, :], state)
        #         outputs.append(cell_output)
        # # output shape: (B * T, D)
        # output = tf.reshape(tf.stack(axis=1, values=outputs), [-1, size])
        





        softmax_w = tf.get_variable(
                "softmax_w", [size, vocab_size], dtype=data_type())
        softmax_b = tf.get_variable("softmax_b", [vocab_size], dtype=data_type())
        logits = tf.matmul(output, softmax_w) + softmax_b

        # Reshape logits to be 3-D tensor for sequence loss
        logits = tf.reshape(logits, [batch_size, num_steps, vocab_size])

        # use the contrib sequence loss and average over the batches
        loss = tf.contrib.seq2seq.sequence_loss(
                logits,
                input_.targets,
                tf.ones([batch_size, num_steps], dtype=data_type()),
                average_across_timesteps=False,
                average_across_batch=True
        )
        # loss shape: (T)


        correct_prediction = tf.equal(
            tf.argmax(tf.nn.softmax(logits), 2),
            tf.cast(input_.targets, tf.int64))
        accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))


        # update the cost variables
        # To conclude, loss avg on B, sum on T
        self._cost = cost = tf.reduce_sum(loss)
        self._acc = accuracy
        self._final_state = state

        if not is_training:
            return

        self._lr = tf.Variable(0.0, trainable=False)
        tvars = tf.trainable_variables()
        grads, _ = tf.clip_by_global_norm(tf.gradients(cost, tvars),
                                                                            config.max_grad_norm)
        optimizer = tf.train.GradientDescentOptimizer(self._lr)
        self._train_op = optimizer.apply_gradients(
                zip(grads, tvars),
                global_step=tf.contrib.framework.get_or_create_global_step())

        self._new_lr = tf.placeholder(
                tf.float32, shape=[], name="new_learning_rate")
        self._lr_update = tf.assign(self._lr, self._new_lr)

    def assign_lr(self, session, lr_value):
        session.run(self._lr_update, feed_dict={self._new_lr: lr_value})

    @property
    def input(self):
        return self._input

    @property
    def initial_state(self):
        return self._initial_state

    @property
    def cost(self):
        return self._cost

    @property
    def acc(self):
        return self._acc

    @property
    def final_state(self):
        return self._final_state

    @property
    def lr(self):
        return self._lr

    @property
    def training(self):
        return self._training

    @property
    def train_op(self):
        return self._train_op